'''
yolov8 分类推理
'''
import glob
import os.path
import time

import cv2
import numpy as np
import torch.cuda

from ultralytics import YOLO # v8.1
from line_profiler import LineProfiler # 逐行计算运行时间
ctime = LineProfiler()

def get_yolov8_cls_predict(model_path = 'yolov8n-cls.pt'):
    # model_path = 'yolov8n.pt'
    if isinstance(model_path, YOLO):
        model = model_path
    else:
        model = YOLO(model_path)
    device = '0' if torch.cuda.is_available() else 'cpu'
    model.predict(np.zeros((640,640,3), np.uint8), imgsz=640, device=device, retina_masks=True) # 预热
    # model.predictor(imgsz=320, conf=0.5)
    # @ctime
    def yolov8_cls_predict(img, imgsz=640, realtime_conf_thres=0.35, roi=None, pad=0,
                            show_img=None, thickness_rate=0.001):
        '''
        推理单张图片， 可roi分类
        show_img, out = seg_v8_step1xjg(img0, realtime_conf_thres=0.5, iou_thres=0.45,
                                        roi=fix_bd_step1['xiangjigai'], pad=20, isShowMask=False)
        Args:
            img: 2048*3072*3 原图
            realtime_conf_thres:
            iou_thres:
            roi: [xyxy] 是否roi
            pad: 填充 px
            isShowMask: 是否画mask
            show_img : 可外部传入show_img, 为None则copy原图
            isMask2Xy: mask是否转为轮廓(numpy 1*n*2)，加快速度
            isDraw: 是否画图
        Returns:
            show_img:
            out: [clsid, conf] # [int, float]

        '''
        H, W = img.shape[:2]
        if roi is not None:
            x1, y1, x2, y2 = roi
            roi = [int(max(x1 - pad, 0)), int(max(y1 - pad, 0)), int(min(x2 + pad, W)), int(min(y2 + pad, H))]

        img2 = img[roi[1]:roi[3],roi[0]:roi[2],...] if roi is not None else img
        t0 = time.time()
        results = model.predict(img2, imgsz=imgsz, conf=realtime_conf_thres, device=device)  # retina_masks 高分辨率mask
        print(f'3333333333333333333 {time.time()-t0}')
        result = results[0]  # one img
        probs = result.probs  # Probs object for classification outputs

        # out = [probs.top1] # [int]
        # out = [probs.top1, probs.top1conf.cpu().numpy().item()]  # [clsid, conf]
        out = [model.names[probs.top1], probs.top1conf.cpu().numpy().item()] # [clsname, conf]
        print(out)

        # draw
        if show_img is None:
            show_img = img.copy()
        img_diag = np.sqrt(show_img.shape[0] ** 2 + show_img.shape[1] ** 2)
        print(img_diag)
        cls_ind, cls_conf = out
        px1, py1 = 10, 10
        if roi is not None:
            px1 = roi[0] + px1
            py1 = roi[1] + py1
        print(show_img.shape)
        cv2.putText(show_img, f'{cls_ind} {cls_conf:.2f}', (px1, py1),
                    cv2.FONT_HERSHEY_COMPLEX_SMALL, max(1, int(img_diag * thickness_rate*0.5)), (0, 255, 0),
                    max(1, int(img_diag * 0.0005)))

        if roi is not None:
            cv2.rectangle(show_img, (roi[0], roi[1]), (roi[2], roi[3]), (255, 0, 0), max(1, int(img_diag * thickness_rate)))
        print(f'555555555555555 {time.time() - t0}')
        return show_img, out

    return yolov8_cls_predict



def ttest_get_detect_onepic():
    img = cv2.imread(r"D:\data\231207huoni\trainV8Seg_cable\add_imgs\20240104\Image_20240104150413321.jpg")
    model_path = r"D:\data\240815fubahanxi\trainV8Seg_fubahanxi\models\yolov8sCls_160_fubahanxi3\weights\best.pt"
    detect_onepic = get_yolov8_cls_predict(model_path)
    img_show, out = detect_onepic(img,imgsz=160,roi=[100,100,2500,1500], isShowMask=True, isMask2Xy=True)
    ctime.print_stats()
    cv2.imwrite(r'D:\data\231207huoni\test_data\1.jpg', img_show)
    # print(out)

def ttest_dir(model_path = '', img_glob = '', save_dir = ''):
    # from ultralytics import YOLO
    # test_dir = r"D:\data\231215安全带\trainV8Det_flball_blue\format_data\images\train"
    # # Load a pretrained YOLOv8n model
    # model_path = r"D:\data\231215安全带\trainV8Det_flball_blue\models\train7\weights\best.pt"
    # model = YOLO(model_path)
    # # Run inference on 'bus.jpg' with arguments
    # model.predict(test_dir, save=True, imgsz=640, conf=0.5, show=True)

    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    ls = glob.glob(img_glob)
    cls_predict = get_yolov8_cls_predict(model_path)
    for ind, i in enumerate(ls):
        img = cv2.imread(i)
        # print(i)
        # print(img)
        # try:
        img_show, out = cls_predict(img, 640, show_img=img)
        cv2.imwrite(f'{save_dir}/{ind}.jpg', img_show)
        # except:
        #     print(f'error {i}')

if __name__ == '__main__':
    # ttest_get_detect_onepic()
    root_path = r'D:\DATA\20250519RENBAO\trainV8Pose_people'
    # # root_path = r'/home/ps/zhangxiancai/data/231207huoni/'
    # ttest_dir(model_path = rf"{root_path}/models\yolov8sCls_320_people_20250523113200\weights\best.pt",
    #           img_glob = rf'{root_path}/cls_format_data/test/*/*.jpg',
    #           save_dir = rf"{root_path}/predict")

    ttest_dir(model_path = rf"{root_path}/models\yolov8sCls_320_people_20250603_113242_yolov8_cls__cls_format_data_img100_200_cls7_augmean3_pad\weights\best.pt",
              img_glob = rf'{root_path}\cls_format_data_img100_200_cls7_augmean3_pad\test\ZhCMFangBao/*.jpg',
              save_dir = rf"{root_path}/predict")
